#include <stdio.h>
#include <stdlib.h>
#include <iostream>
#include <string>
#include <vector>
#include <fstream>
#include "matmul.h"
#include <iomanip>
// #include <iterator>

// #define CORRECTNESS

int main(int argc, char const *argv[])
{
    cl_int status = 0;
    printf("\n------------------------------ main() start ------------------------------\n");
    mslite::LoadOpenCL();
    mslite::Enviroment::createEnv(1);

    #ifdef CORRECTNESS 
    std::vector<int> lhsShape = {128, 64}, rhsShape = {64, 128};            // verify correctness
    #else
    std::vector<int> lhsShape = {512, 512}, rhsShape = {512, 512};        // verify performance
    #endif
    std::cout << "lhs: [" << lhsShape[0] << ", " << lhsShape[1] << "], ";
    std::cout << "rhs: [" << rhsShape[0] << ", " << rhsShape[1] << "]" << std::endl;
    mslite::Matmul<int8_t> *int8_op = new mslite::Matmul<int8_t>("matmul_int8_lbriobntnp", lhsShape, rhsShape, 0);
    mslite::Matmul<float>  *fp32_op = new mslite::Matmul<float >("matmul_fp32_lbriobntnp", lhsShape, rhsShape, 0);
    int8_op->setGlobal({static_cast<size_t>(lhsShape[0]), static_cast<size_t>(rhsShape[1])});

    std::vector<std::vector<int8_t>> int8_inputs(2);
    std::vector<std::vector<float >> fp32_inputs(2);
    int8_inputs[0] = mslite::genData(lhsShape[0]*lhsShape[1], static_cast<int8_t>(-16), static_cast<int8_t>(16));
    int8_inputs[1] = mslite::genData(rhsShape[0]*rhsShape[1], static_cast<int8_t>(-16), static_cast<int8_t>(16));
    fp32_inputs[0] = mslite::convertToFloat(int8_inputs[0]);
    fp32_inputs[1] = mslite::convertToFloat(int8_inputs[1]);
    int8_op->setInputs(int8_inputs);
    fp32_op->setInputs(fp32_inputs);

    const std::vector<std::vector<size_t>> localCandidates 
        = {{16, 16}, {16, 8}, {8, 16}, {8, 8}, {8, 4}, {4, 8}, {4, 4}, {4, 2}, {2, 4}, {2, 2}};
    int testRound = 20, discard = 5;
    std::vector<double> int8_times, fp32_times;
    std::cout << "---------------------------------------------------------" << std::endl;
    std::cout << " Local" << "\t\tInt8-Op-Time" << "\tFp32-Op-Time" << "\tFp32/Int8" << std::endl;
    std::cout << "---------------------------------------------------------" << std::endl;
    for (auto candi: localCandidates) {
        int8_op->setLocal(candi);
        fp32_op->setLocal(candi);
        int8_times.clear();
        fp32_times.clear();

        for (int j=0; j<testRound; j++) {
            cl_event profEvent;
            int8_op->run(profEvent);
            int8_times.emplace_back(mslite::profileTime(profEvent, false));
        }

        for (int j=0; j<testRound; j++) {
            cl_event profEvent;
            fp32_op->run(profEvent);
            fp32_times.emplace_back(mslite::profileTime(profEvent, false));
        }

        double int8_avg = mslite::getAvg(int8_times, discard), fp32_avg = mslite::getAvg(fp32_times, discard);

        std::cout << "(" << std::setw(2) << candi[0] << ", " << std::setw(2) << candi[1] << ")\t ";
        std::cout << std::setw(7) << int8_avg << "ms\t " << std::setw(7) << fp32_avg << "ms\t" << std::setw(7) << fp32_avg / int8_avg << "x" << std::endl;
    }
    std::cout << "---------------------------------------------------------" << std::endl;

    #ifdef CORRECTNESS
    std::vector<float > expect = mslite::matmul(fp32_inputs[0], fp32_inputs[1], lhsShape[0], lhsShape[1], rhsShape[1]);
    std::vector<int8_t> int8_out = int8_op->getOutput()[0];
    std::vector<float > fp32_out = fp32_op->getOutput()[0];
    double int8Diff = 0, fp32Diff = 0;
    for (int i=0; i<int8_out.size(); i++) {
        int8Diff = std::max(int8Diff, std::fabs(static_cast<double>(int8_out[i]-static_cast<int>(expect[i])%53)));
        fp32Diff = std::max(fp32Diff, std::fabs(static_cast<double>(static_cast<int>(fp32_out[i])%53)-static_cast<int>(expect[i])%53));
        // if (i < 10) {
            // std::cout << static_cast<int>(int8_out[i]) << " " << fp32_out[i] << " " << expect[i] << std::endl;
        // }
    }
    std::cout << "Max Difference: " << int8Diff << " " << fp32Diff << std::endl;
    #endif

    status = mslite::Enviroment::releaseEnv();
    printf("------------------------------- main() over -------------------------------\n");
    return 0;
}